# this file plan a traj for lateral control within one cycle time based on sample alg
# the alg processing is as follows:
# 1. sample longitude velocity and lateral position
# 2. merge longitude and lateral traj
# 3. select optimize traj
# (1) filter invalid acceleration and yawrate
# (2) collision detecte
# (3) calc cost function return the minimun one

import matplotlib.pyplot as plt
import numpy as np
from sympy import *
import scipy.integrate as intg
import math


# struct to sort coeff of a traj about sample time and lateral and longitude
class coeff_strut:
    def __init__(self, lateral_coeff, longitude_coeff, pred_t):
        self.lateral_coeff = lateral_coeff
        self.longitude_coeff = longitude_coeff
        self.predict_t = pred_t

    def get_lateral_coeff(self):
        return self.lateral_coeff

    def get_longitude_coeff(self):
        return self.longitude_coeff

    def get_predt(self):
        return self.predict_t


class TrajCoeff:
    def __init__(self, coeff, sampletime):
        self.coeff = coeff
        self.sampletime = sampletime

    def getCoeff(self):
        return self.coeff

    def getSampletime(self):
        return self.sampletime


# struct to store vehicle kinetics state
class vehicle_state:
    def __init__(self, x, y, vx, vy, ax, ay):
        self.x = x
        self.y = y
        self.vx = vx
        self.vy = vy
        self.ax = ax
        self.ay = ay

    def get_x(self):
        return self.x

    def get_y(self):
        return self.y

    def get_vx(self):
        return self.vx

    def get_vy(self):
        return self.vy

    def get_ax(self):
        return self.ax

    def get_ay(self):
        return self.ay


# this fun is used to transfer the ego kinetics state to reference line
# return : map{
# "vx2ref":ego_vx2refline,
# "ax2ref":ego_ax2refline,
# "vy2ref":ego_vy2refline,
# "ay2ref":ego_ay2refline
# }
def transfer2refline(target, ego):
    delta_x = target.get_x() - ego.get_x()
    delta_y = target.get_y() - ego.get_y()
    thirdline = (delta_x ** 2 + delta_y ** 2) ** 0.5
    cos_theta = delta_y / thirdline
    sin_theta = delta_x / thirdline
    # lateral velocity and accelerate
    ego_v2refline = ego.get_vx() * cos_theta - ego.get_vy() * sin_theta
    ego_a2refline = ego.get_ax() * cos_theta - ego.get_ay() * sin_theta
    # longitude velocity and accelerate
    ego_vOnrefline = ego.get_vy() * cos_theta + ego.get_vx() * sin_theta
    ego_aOnrefline = ego.get_ay() * cos_theta + ego.get_ax() * sin_theta

    return {"v2ref": ego_v2refline, "a2ref": ego_a2refline, "vOnref": ego_vOnrefline, "aOnref": ego_aOnrefline}


# this fun is used to sample lateral postion
def lateral_sampling(steptime_array, origion_state):
    d_array = np.arange(0.5, 4.5, 0.5)
    # distance, velocity, acceleration
    finalstate_array = []
    for i in d_array:
        finalstate_array.append([i, 0, 0])

    lateral_coeff_map = {}
    for t in steptime_array:
        coeff_list = []
        for final_state in finalstate_array:
            O = origion_state + final_state
            M = np.array([
                [1, 0, 0, 0, 0, 0],
                [0, 1, 0, 0, 0, 0],
                [0, 0, 2, 0, 0, 0],
                [1, t, t ** 2, t ** 3, t ** 4, t ** 5],
                [0, 1, 2 * t, 3 * t ** 2, 4 * t ** 3, 5 * t ** 4],
                [0, 0, 2, 6 * t, 12 * t ** 2, 20 * t ** 3]
            ])
            O = np.array([O]).T
            t_co = np.dot(np.linalg.inv(M), O)
            t_co = np.flipud(t_co).flatten()
            t_fun = np.poly1d(t_co.astype(float))
            coeff_list.append(t_fun)
        lateral_coeff_map[t] = coeff_list
    return lateral_coeff_map


# this fun is used to sample longitude velocity
def longitude_sampling(steptime_array, origion_state):
    step_num = 5
    v_array = np.linspace(15, 14, step_num)
    a_array = np.linspace(0, 1, step_num)
    finalstate_array = []
    for i in range(step_num):
        finalstate_array.append([v_array[i], a_array[i]])

    longitude_coeff_map = {}
    for t in steptime_array:
        coeff_list = []
        for final_longitude_state in finalstate_array:
            O = [0] + origion_state + final_longitude_state
            M = [
                [1, 0, 0, 0, 0],  # 初始位移
                [0, 1, 0, 0, 0],  # 初始速度
                [0, 0, 2, 0, 0],  # 初始加速度
                [0, 1, 2 * t, 3 * t ** 2, 4 * t ** 3],
                [0, 0, 2, 6 * t, 12 * t ** 2]
            ]
            O = np.array([O]).T
            t_co = np.dot(np.linalg.inv(M), O)
            t_co = np.flipud(t_co).flatten()
            t_fun = np.poly1d(t_co.astype(float))
            coeff_list.append(t_fun)
            # store formula not coeff
        longitude_coeff_map[t] = coeff_list
    return longitude_coeff_map


# merge lateral and longitude traj
def merge_traj(lateral_traj_map, longitude_traj_map):
    merged_traj = []
    for sampletime in lateral_traj_map.keys():
        for lat in lateral_traj_map[sampletime]:
            for long in longitude_traj_map[sampletime]:
                merged_traj.append({"predtime": sampletime, "latTraj": lat, "longTraj": long})
    return merged_traj


# filter invalid lateral and longitude acceleration
def filter_traj(Descartes_traj_set):
    for index in range(len(Descartes_traj_set) - 1, 0, -1):
        traj = Descartes_traj_set[index]
        sampleTime = np.arange(0, traj["predtime"], 0.02)

        # lat_coeff=traj["latTraj"][::-1]
        # lat_coeff=[x for i in lat_coeff for x in i]
        # diff_lat_coeff=np.polyder(np.polyder(lat_coeff))
        diff_lat = np.polyder(np.polyder(traj["latTraj"]))
        lat_acceleration = np.polyval(diff_lat, sampleTime)
        if (lat_acceleration.max() > 2 or lat_acceleration.min() < -2):
            Descartes_traj_set.pop(index)
            continue

        long_coeff = traj["longTraj"]
        diff_long = np.polyder(np.polyder(long_coeff))
        long_acceleration = np.polyval(diff_long, sampleTime)
        if (long_acceleration.max() > 2 or long_acceleration.min() < -2):
            Descartes_traj_set.pop(index)
            # break


# this fun is used to filter trajs which will collide with obj (static obj)
# using SL graphic
def collision_detect(Descartes_traj_set, target_state, ego_state):
    smin = sqrt((target_state.get_y() - ego_state.get_y()) ** 2 + (target_state.get_x() - ego_state.get_x()) ** 2)
    smax = smin + 4.5
    lmin = 1
    lmax = -1

    for index in range(len(Descartes_traj_set) - 1, -1, -1):
        traj = Descartes_traj_set[index]
        for sampleTime in np.arange(0, 5, 0.02):
            long_coeff = traj["longTraj"]
            # 积分得到S，三角函数得到d
            s = intg.quad(long_coeff, 0, sampleTime)[0]
            lat_coeff = traj["latTraj"]
            d = lat_coeff(sampleTime)

            if (s < smax and s > smin and d < lmin and d > lmax):
                Descartes_traj_set.pop(index)
                break


def optimal_traj(Descartes_traj_set):
    min_cost = 9999
    optimized_traj = {}
    lat_weight = 0.5
    long_weight = 0.5
    for traj in Descartes_traj_set:
        sampleTime = np.arange(0, traj["predtime"], 0.02)

        lat_coeff = traj["latTraj"]
        diff_lat_coeff = np.polyder(np.polyder(np.polyder(lat_coeff)))
        lat_jerk = np.polyval(diff_lat_coeff, sampleTime).sum()

        long_coeff = traj["longTraj"]
        diff_long_coeff = np.polyder(np.polyder(long_coeff))
        long_jerk = np.polyval(diff_long_coeff, sampleTime).sum()

        temp_cost = lat_weight * lat_jerk + long_weight * long_jerk < min_cost
        if (temp_cost < min_cost):
            min = temp_cost
            optimized_traj = traj
    return optimized_traj


def calc_refline(target_state, ego_state):
    target_x = target_state.get_x()
    target_y = target_state.get_y()
    ego_x = ego_state.get_x()
    ego_y = ego_state.get_y()

    k = (target_y - ego_y) / (target_x - ego_x)
    b = target_y - k * target_x
    return {"k": k, "b": b}


def getLocation_in_global(sample_time, traj, refline, ego_state):
    # long_coeff = optimized_traj["longTraj"][::-1]
    # long_coeff = np.array([coeff for i in long_coeff for coeff in i])
    # # 积分得到S，三角函数得到d
    # func = np.poly1d(long_coeff.astype(float))
    long_fun = traj["longTraj"]
    s = intg.quad(long_fun, 0, sample_time)[0]

    # lat_coeff = optimized_traj["latTraj"][::-1]
    # lat_coeff = [coeff for i in lat_coeff for coeff in i]
    # d=np.polyval(lat_coeff,sample_time)
    lat_fun = traj["latTraj"]
    L = lat_fun(sample_time)

    sx = math.cos(math.atan(refline["k"])) * s
    sy = math.sin(math.atan(refline["k"])) * s

    sOnref_x = ego_state.get_x() + sx
    sOnref_y = ego_state.get_y() + sy

    dx = math.sin(math.atan(refline["k"])) * L
    dy = math.cos(math.atan(refline["k"])) * L

    global_x = sOnref_x - dx
    global_y = sOnref_y + dy

    return [global_x, global_y]


if __name__ == '__main__':

    # obstacle property
    width = 2
    lenght = 4.5
    # target location
    target_state = vehicle_state(20.5, 60, 0, 0, 0, 0)

    # ego original state
    ego_state = vehicle_state(20, 0, 0, 15, 0, 0.5)

    # trans ego state to reference line
    origion_state = transfer2refline(target_state, ego_state)

    # predict time
    predtime_ar = np.arange(0.5, 5, 0.5)

    refline = calc_refline(target_state, ego_state)
    # lateral planning
    lateral_origion_state = [0, origion_state["v2ref"], origion_state["a2ref"]]
    lateral_coeff_map = lateral_sampling(predtime_ar, lateral_origion_state)
    plt.figure(1)
    for predt in lateral_coeff_map.keys():
        for traj in lateral_coeff_map[predt]:
            x = np.linspace(0, predt, 30)
            y = traj(x)
            plt.plot(x, y)

    # longitude planning
    longitude_origion_state = [origion_state["vOnref"], origion_state["aOnref"]]
    longitude_coeff_map = longitude_sampling(predtime_ar, longitude_origion_state)
    plt.figure(2)
    for predt in longitude_coeff_map.keys():
        for traj in longitude_coeff_map[predt]:
            x = np.linspace(0, predt, 30)
            y = np.polyder(traj)(x)
            plt.plot(x, y)

    # merge lateral and longitude traj
    merged_traj_list = merge_traj(lateral_coeff_map, longitude_coeff_map)
    print("traj_num = " + str(len(merged_traj_list)))
    filter_traj(merged_traj_list)
    print("traj_num after filter: " + str(len(merged_traj_list)))
    collision_detect(merged_traj_list, target_state, ego_state)
    print("traj_num after collision detect: " + str(len(merged_traj_list)))
    optimized_traj = optimal_traj(merged_traj_list)

    plt.figure(3)
    plt.scatter(target_state.get_x(), target_state.get_y())
    plt.scatter(target_state.get_x() - 1, target_state.get_y())
    plt.scatter(target_state.get_x() + 1, target_state.get_y())
    plt.scatter(ego_state.get_x(), ego_state.get_y())
    for tTraj in merged_traj_list:
        traj_point_x = []
        traj_point_y = []
        for sampleTime in np.arange(0, tTraj["predtime"], 0.2):
            temp_point = getLocation_in_global(sampleTime, tTraj, refline, ego_state)
            traj_point_x.append(temp_point[0])
            traj_point_y.append(temp_point[1])
        plt.plot(traj_point_x, traj_point_y)
    plt.xlim(-10, 50)
    plt.ylim(-10, 200)

    plt.figure(4)
    plt.scatter(target_state.get_x(), target_state.get_y())
    plt.scatter(target_state.get_x() - 1, target_state.get_y())
    plt.scatter(target_state.get_x() + 1, target_state.get_y())
    plt.scatter(ego_state.get_x(), ego_state.get_y())
    traj_point_x = []
    traj_point_y = []
    for sampleTime in np.arange(0, optimized_traj["predtime"], 0.2):
        temp_point = getLocation_in_global(sampleTime, optimized_traj, refline, ego_state)
        traj_point_x.append(temp_point[0])
        traj_point_y.append(temp_point[1])
    plt.plot(traj_point_x, traj_point_y)

    plt.xlim(-10, 50)
    plt.ylim(-10, 200)
    plt.show()
    # draw env
    # plt.figure("env")
    # plt.scatter(target_state.get_x(),target_state.get_y())
    # plt.scatter(target_state.get_x()-1,target_state.get_y())
    # plt.scatter(target_state.get_x()+1,target_state.get_y())
    # plt.scatter(ego_state.get_x(),ego_state.get_y())
    #
    # for traj in merged_traj_list:
    #     traj_point_x=[]
    #     traj_point_y=[]
    #     for sampleTime in np.arange(0,1,0.2):
    #         temp_point=getLocation_in_global(sampleTime,traj, refline, ego_state)
    #         traj_point_x.append(temp_point[0])
    #         traj_point_y.append(temp_point[1])
    #     plt.plot(traj_point_x,traj_point_y)
    #
    # plt.xlim(-10,50)
    # plt.ylim(-10,50)
    # plt.show()

